#include <pybind11/chrono.h>
#include <pybind11/eigen.h>
#include <pybind11/functional.h>
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/stl_bind.h>

#include <iostream>

#include "common.h"
#include "ouster/client.h"
#include "ouster/impl/profile_extension.h"
#include "ouster/lidar_scan.h"
#include "ouster/osf/basics.h"
#include "ouster/osf/meta_extrinsics.h"
#include "ouster/osf/meta_lidar_sensor.h"
#include "ouster/osf/meta_streaming_info.h"
#include "ouster/osf/metadata.h"
#include "ouster/osf/operations.h"
#include "ouster/osf/reader.h"
#include "ouster/osf/stream_lidar_scan.h"
#include "ouster/osf/writer.h"

namespace py = pybind11;

using namespace ouster;

inline std::vector<uint8_t> getvector(py::buffer& buf) {
    auto info = buf.request();
    if (info.format != py::format_descriptor<uint8_t>::format()) {
        throw std::invalid_argument(
            "Incompatible argument: expected a bytearray");
    }
    if (info.ndim != 1) {
        throw std::invalid_argument(
            "Incompatible argument: expect number of dimensions 1");
    }
    return {(uint8_t*)info.ptr, (uint8_t*)info.ptr + info.size};
}

/*
 * NOTE: Copied from _client.cpp
 * Map a dtype to a channel field type
 */
static sensor::ChanFieldType field_type_of_dtype(const py::dtype& dt) {
    if (dt.is(py::dtype::of<uint8_t>()))
        return sensor::ChanFieldType::UINT8;
    else if (dt.is(py::dtype::of<uint16_t>()))
        return sensor::ChanFieldType::UINT16;
    else if (dt.is(py::dtype::of<uint32_t>()))
        return sensor::ChanFieldType::UINT32;
    else if (dt.is(py::dtype::of<uint64_t>()))
        return sensor::ChanFieldType::UINT64;
    else
        throw std::invalid_argument("Invalid dtype for a channel field");
}

void init_osf(py::module& m, py::module&) {
    m.doc() = R"doc(
Ouster OSF Python API generated by pybind11.

This module is generated from the C++ code and provides native functions
to work with OSF files.
)doc";

    m.def("dump_metadata", &ouster::osf::dump_metadata, R"doc(
        Dump OSF metadata/session info in JSON format. (aka osf-metadata)

        :file: OSF file path
        :returns: JSON formatted string of OSF metadata + header info
    )doc",
          py::arg("file"), py::arg("full") = true);

    m.def("parse_and_print", &ouster::osf::parse_and_print, R"doc(
        Parse OSF file and print messages types, timestamps and counts to
        stdout.

        :file: OSF file path (v1/v2)
    )doc",
          py::arg("file"), py::arg("with_decoding") = false);

    m.def("backup_osf_file_metablob", &ouster::osf::backup_osf_file_metablob,
          R"doc(
         Backup the metadata blob in an OSF file.

        :file: OSF file path (v1/v2)
        :backup_file_name: Backup path
    )doc",
          py::arg("file"), py::arg("backup_file_name"));

    m.def("restore_osf_file_metablob", &ouster::osf::restore_osf_file_metablob,
          R"doc(
        Restore an OSF metadata blob from a backup file.

        :file: OSF file path (v1/v2)
        :backup_file_name: The backup to use
    )doc",
          py::arg("file"), py::arg("backup_file_name"));

    m.def("osf_file_modify_metadata", &ouster::osf::osf_file_modify_metadata,
          R"doc(
        Modify an OSF files sensor_info metadata.

        :file_name: The OSF file to modify.
        :new_metadata: Array containing sensor infos to write to the file.
        :returns: The number of the bytes written to the OSF file.
    )doc",
          py::arg("file_name"), py::arg("new_metadata"));

    // Reader
    py::class_<osf::Reader>(m, "Reader", R"(
        Reader is a main entry point to get any info out of OSF file.
    )")
        .def(py::init<std::string>(), py::arg("file"))
        .def_property_readonly("metadata_id", &osf::Reader::metadata_id, R"(
            Data id string
        )")
        .def_property_readonly(
            "start_ts",
            [](const osf::Reader& r) { return r.start_ts().count(); }, R"(
                Start timestamp (ns) - the lowest message timestamp present in the file
        )")
        .def_property_readonly(
            "end_ts", [](const osf::Reader& r) { return r.end_ts().count(); },
            R"(
                End timestamp (ns) - the highest message timestamp present in the file
        )")
        .def_property_readonly("meta_store", &osf::Reader::meta_store, R"(
                Returns the metadata store that gives an access to all
                *metadata entries* in the file.
        )")
        .def_property_readonly(
            "has_stream_info", &osf::Reader::has_stream_info,
            "Whether ``StreamingInfo`` metadata is available (i.e. reading "
            "messages by timestamp and streams can be performed)")
        .def_property_readonly(
            "has_message_idx", &osf::Reader::has_message_idx,
            "Whether OSF contains the message counts that are needed for "
            "``ts_by_message_idx()`` (message counts was added a bit later to "
            "the OSF core, so this function will be obsolete over time)")
        .def_property_readonly(
            "has_timestamp_idx", &osf::Reader::has_timestamp_idx,
            "Whether OSF contains the message timestamp index in the metadata "
            "necessary to quickly collate and jump to a specific message time.")
        .def(
            "messages",
            [](osf::Reader& r) {
                return py::make_iterator(r.messages().begin(),
                                         r.messages().end());
            },
            py::keep_alive<0, 1>(), R"(
                Creates an iterator to reads messages in default ``STREAMING`` layout.
            )")
        .def(
            "messages",
            [](osf::Reader& reader, uint64_t start_ts, uint64_t end_ts) {
                auto msgs =
                    reader.messages(osf::ts_t{start_ts}, osf::ts_t{end_ts});
                return py::make_iterator(msgs.begin(), msgs.end());
            },
            py::keep_alive<0, 1>(), py::arg("start_ts"), py::arg("end_ts"),
            R"(
                    Read `messages` in ``[start_ts, end_ts]`` timestamp range
                    (inclusive)
                )")
        .def(
            "messages",
            [](osf::Reader& reader, std::vector<uint32_t> stream_ids) {
                auto msgs = reader.messages(stream_ids);
                return py::make_iterator(msgs.begin(), msgs.end());
            },
            py::keep_alive<0, 1>(),
            py::arg("stream_ids") = std::vector<uint32_t>{},
            R"(
                    Read `messages` from only specified ``[<stream_ids>]`` list
                )")
        .def(
            "messages",
            [](osf::Reader& reader, std::vector<uint32_t> stream_ids,
               uint64_t start_ts, uint64_t end_ts) {
                auto msgs = reader.messages(stream_ids, osf::ts_t{start_ts},
                                            osf::ts_t{end_ts});
                return py::make_iterator(msgs.begin(), msgs.end());
            },
            py::keep_alive<0, 1>(), py::arg("stream_ids"), py::arg("start_ts"),
            py::arg("end_ts"),
            R"(
                    Read `messages` in ``[start_ts, end_ts]`` timestamp range (inclusive) of a
                    specified ``<stream_ids>`` list
                )")
        .def(
            "ts_by_message_idx",
            [](osf::Reader& reader, uint32_t stream_id,
               uint32_t message_idx) -> py::object {
                auto ts = reader.ts_by_message_idx(stream_id, message_idx);
                if (ts) {
                    return py::int_(ts->count());
                }
                return py::none();
            },
            py::arg("stream_id"), py::arg("message_idx"),
            R"(
                    Find the timestamp of the message by its index and stream_id.

                    Requires the OSF with message_counts inside, i.e. has_message_idx()
                    is ``True``, otherwise return value is always None.
                )")
        .def(
            "chunks",
            [](osf::Reader& r) {
                return py::make_iterator(r.chunks().begin(), r.chunks().end());
            },
            py::keep_alive<0, 1>(), R"(
                Creates an iterator to reads chunks as they appear in a file.
            )");

    // MessageRef
    py::class_<osf::MessageRef>(m, "MessageRef", R"(
        Thin `message` wrapper for underlying `StampedMessage` object.
        
        Provides the ``decode()`` function that resolves the underlying message
        bytes into the object according to the `message` type.

        Underlying message memory is not copied and reference to mmap'ed file
        object until the ``decode()`` is called.
    )")
        .def_property_readonly(
            "id", &osf::MessageRef::id,
            "Message id which is a ``stream_id`` and point to the "
            "`metadata entry` that describes the stream")
        .def_property_readonly(
            "ts", [](const osf::MessageRef& msg) { return msg.ts().count(); },
            "Message timestamp (ns)")
        .def_property_readonly("buffer", &osf::MessageRef::buffer,
                               "Returns encoded message byte array")
        .def("__repr__", &osf::MessageRef::to_string)
        .def("__str__", &osf::MessageRef::to_string)
        .def(
            "of",
            [](const osf::MessageRef& msg, py::object msg_stream) {
                if (py::hasattr(msg_stream, "type_id")) {
                    std::string type_id = py::cast<std::string>(
                        py::getattr(msg_stream, "type_id"));
                    return msg.is(type_id);
                }
                return false;
            },
            py::arg("msg_stream"),
            "Checks whether the message belongs to the ``msg_stream`` type")
        .def(
            "decode",
            [](const osf::MessageRef& msg,
               const std::vector<std::string>& fields) -> py::object {
                if (msg.is<osf::LidarScanStream>()) {
                    auto decoded_obj =
                        msg.decode_msg<osf::LidarScanStream>(fields);
                    return py::cast(*decoded_obj);
                }
                // TODO[pb]: Add dynamic check for Stream decoding functions ...
                return py::none();
            },
            py::arg("fields") = std::vector<std::string>(),
            py::return_value_policy::move,
            R"(
            Decodes the underlying object and returns it.

            Currently supports only LidarScans
        )");

    // MetadataStore
    py::class_<osf::MetadataStore>(m, "MetadataStore", R"(
        Stores `metadata entries` of the file.

        One of information about available sensors, it's configuration,
        available streams, chunks layout method, etc.
    )")
        .def(py::init<>())
        .def("__len__", &osf::MetadataStore::size,
             "Number of `metadata entries` in the file")
        .def(
            "__iter__",
            [](const osf::MetadataStore& m) {
                return py::make_key_iterator(m.entries().begin(),
                                             m.entries().end());
            },
            py::keep_alive<0, 1>(), "Creates an iterator to get metadata id's")
        .def(
            "__getitem__",
            [](const osf::MetadataStore& m, osf::MetadataStore::key_type k) {
                auto entries = m.entries();
                auto it = entries.find(k);
                if (it == entries.end()) {
                    throw py::key_error();
                }
                return it->second.get();
            },
            py::arg("meta_id"), py::return_value_policy::reference,
            "Get `metadata entry` by id")
        .def(
            "items",
            [](const osf::MetadataStore& m) {
                return py::make_iterator(m.entries().begin(),
                                         m.entries().end());
            },
            py::keep_alive<0, 1>(), "Key/Value pairs of `metadata entries`")
        .def(
            "find",
            [](const osf::MetadataStore& m, py::object mo) {
                std::map<uint32_t, std::shared_ptr<osf::MetadataEntry>> res;
                if (py::hasattr(mo, "type_id")) {
                    std::string type_id =
                        py::cast<std::string>(py::getattr(mo, "type_id"));
                    auto& entries = m.entries();
                    auto it = entries.begin();
                    while (it != entries.end()) {
                        if (type_id == it->second->type()) {
                            res.insert(std::make_pair(it->first, it->second));
                        }
                        ++it;
                    }
                }
                return res;
            },
            py::arg("meta_type"),
            "Get all `metadata entries` of the specified ``meta_type``")
        .def(
            "get",
            [](const osf::MetadataStore& m,
               py::object mo) -> std::shared_ptr<osf::MetadataEntry> {
                if (py::hasattr(mo, "type_id")) {
                    std::string type_id =
                        py::cast<std::string>(py::getattr(mo, "type_id"));
                    auto& entries = m.entries();
                    auto it = entries.begin();
                    while (it != entries.end()) {
                        if (type_id == it->second->type()) {
                            return it->second;
                        }
                        ++it;
                    }
                }
                return nullptr;
            },
            py::arg("meta_type"),
            "Get the first `metadata entry` of the specified ``meta_type``");

    // MetadataEntry --- / ------ <-- trampoline --- / ---
    class PyMetadataEntry : public osf::MetadataEntry {
       public:
        using osf::MetadataEntry::MetadataEntry;

        std::string type() const override {
            PYBIND11_OVERLOAD_PURE(std::string, osf::MetadataEntry, type);
        }

        std::string static_type() const override {
            PYBIND11_OVERLOAD_PURE(std::string, osf::MetadataEntry,
                                   static_type);
        }

        std::unique_ptr<osf::MetadataEntry> clone() const override {
            // NOTE[pb]: Not used here, but needed to make not an abstract class
            return nullptr;
        }

        std::vector<uint8_t> buffer() const override {
            PYBIND11_OVERLOAD_PURE(std::vector<uint8_t>, osf::MetadataEntry,
                                   buffer);
        }

        std::string repr() const override {
            PYBIND11_OVERLOAD(std::string, osf::MetadataEntry, repr);
        }

        std::string to_string() const override {
            PYBIND11_OVERLOAD(std::string, osf::MetadataEntry, to_string);
        }
    };

    // MetadataEntry
    py::class_<osf::MetadataEntry, PyMetadataEntry,
               std::shared_ptr<osf::MetadataEntry>>(m, "MetadataEntry", R"(
        Single OSF `metadata entry`.

        It's typed and has corresponding encoding/decoding functions to the
        underlying bytes representation (``buffer()``/``from_buffer()``)
    )")
        .def(py::init<>())
        .def_property_readonly("type", &osf::MetadataEntry::type,
                               "Type of the metadata entry (use this)")
        .def_property_readonly("static_type", &osf::MetadataEntry::static_type,
                               "Static type, C++ compile time (in Python use "
                               "``type_id`` of concrete type objects instead)")
        .def_property_readonly("id", &osf::MetadataEntry::id,
                               "Id of the metadata entry (unique for a file)")
        .def_property_readonly(
            "buffer", &osf::MetadataEntry::buffer,
            "Encodes (serialize) metadata entry to a stored byte array")
        .def_static(
            "from_buffer",
            [](const std::vector<uint8_t>& buf, const std::string& type_str)
                -> std::shared_ptr<osf::MetadataEntry> {
                return osf::MetadataEntry::from_buffer(buf, type_str);
            },
            py::arg("buf"), py::arg("type_str"),
            "Decodes (deserialize) metadata entry buffer to a typed object")
        .def(
            "of",
            [](const osf::MetadataEntry* m, py::object meta_obj_type) {
                if (py::hasattr(meta_obj_type, "type_id")) {
                    std::string type_id = py::cast<std::string>(
                        py::getattr(meta_obj_type, "type_id"));
                    return (type_id == m->type());
                }
                return false;
            },
            py::arg("meta_obj_type"), R"(
                 Checks whether metadata entry is of particular type

                 It's just:
                 ``self.type == meta_obj_type.type_id``
            )")
        // .def("__repr__", &osf::MetadataEntry::repr)
        .def("__str__", &osf::MetadataEntry::to_string);

    // MetadataEntryRef
    py::class_<osf::MetadataEntryRef, osf::MetadataEntry,
               std::shared_ptr<osf::MetadataEntryRef>>(m, "MetadataEntryRef",
                                                       R"(
        MetadataEntryRef

    )")
        .def_property_readonly_static("type_id", [](py::object) {
            return osf::metadata_type<osf::MetadataEntryRef>();
        });

    // LidarSensor
    py::class_<osf::LidarSensor, osf::MetadataEntry,
               std::shared_ptr<osf::LidarSensor>>(m, "LidarSensor", R"(
        Ouster Lidar Sensor metadata with sensor intrinsics (i.e. SensorInfo/Metadata)

        ``type_id`` static property is a `LidarSensor` metadata type identifier.
    )")
        .def(py::init<sensor::sensor_info>(),
             "Create from ``SensorInfo`` object")
        .def(py::init<std::string>(), py::arg("metadata_json"),
             "Create from ``metadata_json`` string representation")
        .def_property_readonly("info", &osf::LidarSensor::info,
                               "SensorInfo stored",
                               py::return_value_policy::copy)
        .def_property_readonly("metadata", &osf::LidarSensor::metadata,
                               "metadata_json string stored")
        .def_property_readonly_static("type_id", [](py::object) {
            return osf::metadata_type<osf::LidarSensor>();
        });

    // LidarScanStreamMeta
    py::class_<osf::LidarScanStreamMeta, osf::MetadataEntry,
               std::shared_ptr<osf::LidarScanStreamMeta>>(m,
                                                          "LidarScanStreamMeta")
        .def_property_readonly_static(
            "type_id",
            [](py::object) {
                return osf::metadata_type<osf::LidarScanStreamMeta>();
            })
        .def_property_readonly("sensor_meta_id",
                               &osf::LidarScanStreamMeta::sensor_meta_id);

    // LidarScanStream
    py::class_<osf::LidarScanStream>(m, "LidarScanStream", R"(
        `Stream` of ``LidarScan`` objects from a sensor.

        ``type_id`` static property is a `LidarScanStream` underlying metadata
        type identifier.
    )")
        .def_property_readonly_static(
            "type_id",
            [](py::object) {
                return osf::metadata_type<osf::LidarScanStream::meta_type>();
            })
        .def_property_readonly("meta", &osf::LidarScanStream::meta,
                               "`metadata entry` to store `LidarScanStream` "
                               "metadata in an OSF file");

    // StreamStats
    py::class_<osf::StreamStats>(m, "StreamStats", R"(
        Statistics of a stream in ``STREAMING`` layout OSF files.
    )")
        .def_readonly("stream_id", &osf::StreamStats::stream_id,
                      "Id of a stream")
        .def_property_readonly(
            "start_ts",
            [](const osf::StreamStats& ss) { return ss.start_ts.count(); },
            "Lowest timestamp (ns) of the stream messages")
        .def_property_readonly(
            "end_ts",
            [](const osf::StreamStats& ss) { return ss.end_ts.count(); },
            "Highest timestamp (ns) of the stream messages")
        .def_property_readonly(
            "receive_timestamps",
            [](const osf::StreamStats& self) {
                return py::array(
                    py::dtype::of<uint64_t>(), self.receive_timestamps.size(),
                    self.receive_timestamps.data(), py::cast(self));
            },
            "Receive timestamps of each message in the stream.")
        .def_property_readonly(
            "sensor_timestamps",
            [](const osf::StreamStats& self) {
                return py::array(py::dtype::of<uint64_t>(),
                                 self.sensor_timestamps.size(),
                                 self.sensor_timestamps.data(), py::cast(self));
            },
            "Sensor timestamps of each message in the stream.")
        .def_readonly("message_count", &osf::StreamStats::message_count,
                      "Number of messages in a stream")
        .def_readonly("message_avg_size", &osf::StreamStats::message_avg_size,
                      "Average size (bytes) of a message in a stream");

    // StreamingInfo
    py::class_<osf::StreamingInfo, osf::MetadataEntry,
               std::shared_ptr<osf::StreamingInfo>>(m, "StreamingInfo", R"(
        Metadata entry that appears in ``STREAMING`` layout OSF files.

        Establishes the `chunk` -> `stream_id` map as well as providing the
        statistics per stream.

        ``type_id`` static property is a `StreamingInfo` metadata type identifier.
    )")
        .def_property_readonly_static(
            "type_id",
            [](py::object) { return osf::metadata_type<osf::StreamingInfo>(); })
        .def_property_readonly(
            "chunks_info",
            [](osf::StreamingInfo& si) {
                return py::make_iterator(si.chunks_info().begin(),
                                         si.chunks_info().end());
            },
            py::keep_alive<0, 1>(),
            "Maps `chunk` to `stream_id` by chunk offset")
        .def_property_readonly(
            "stream_stats",
            [](osf::StreamingInfo& si) {
                return py::make_iterator(si.stream_stats().begin(),
                                         si.stream_stats().end());
            },
            py::keep_alive<0, 1>(), "Statistics of messages in per stream");

    // Extrinsics
    py::class_<osf::Extrinsics, osf::MetadataEntry,
               std::shared_ptr<osf::Extrinsics>>(m, "Extrinsics", R"(
        Extrinsics transform of a sensor/object referred by ``ref_meta_id``.

        ``type_id`` static property is a ``Extrinsics`` metadata type identifier.
    )")
        .def(py::init<mat4d, uint32_t, std::string&>(), py::arg("extrinsics"),
             py::arg("ref_meta_id") = 0, py::arg("name") = "",
             "Create Extrinsics object")
        .def_property_readonly("extrinsics", &osf::Extrinsics::extrinsics,
                               "Extrisnics homogeneous 4x4 matrix")
        .def_property_readonly("ref_meta_id", &osf::Extrinsics::ref_meta_id,
                               "reference to the metadata entry id of an "
                               "object which extrisnics is it")
        .def_property_readonly("name", &osf::Extrinsics::name,
                               "name of the Extrinsics object (optional)")
        .def_property_readonly_static("type_id", [](py::object) {
            return osf::metadata_type<osf::Extrinsics>();
        });

    // Writer
    py::class_<osf::Writer>(m, "Writer", R"(
        Simple writer interface for OSF file

        All jobs are done with ``MetadataStore`` for adding `metadata entries`
        and stream interfaces that encodes messages and passes them to internal
        chunks writer.
    )")
        .def(py::init<std::string, uint32_t>(), py::arg("file_name"),
             py::arg("chunk_size") = 0, R"(
                 Creates a `Writer` with specified ``chunk_size``.

                 Default ``chunk_size`` is ``2 MB``.
        )")
        .def(py::init([](const std::string& filename,
                         const sensor::sensor_info& info,
                         const std::vector<std::string>& fields_to_write,
                         uint32_t chunk_size) {
                 return new osf::Writer(filename, info, fields_to_write,
                                        chunk_size);
             }),
             py::arg("filename"), py::arg("info"),
             py::arg("fields_to_write") = std::vector<std::string>{},
             py::arg("chunk_size") = 0,
             R"(
            Creates a `Writer` with deafault ``STREAMING`` layout chunks writer.
            
            Using default ``chunk_size`` of ``2MB``.

            Args:
                filename (str): The filename to output to.
                info (sensor_info): The sensor info vector to use for a multi stream OSF
                    file.
                chunk_size (int): The chunk size in bytes to use for the OSF file. This arg
                    is optional, and if not provided the default value of 2MB
                    is used. If the current chunk being written exceeds the
                    chunk_size, a new chunk will be started on the next call to
                    `save`. This allows an application to tune the number of
                    messages (e.g. lidar scans) per chunk, which affects the
                    granularity of the message index stored in the
                    StreamingInfo in the file metadata. A smaller chunk_size
                    means more messages are indexed and a larger number of
                    index entries. A more granular index allows for more
                    precise seeking at the slight expense of a larger file.
                fields_to_write (List[str]): The fields from scans to
                    actually save into the OSF. If not provided uses the fields from 
                    the first saved lidar scan for each stream. This parameter is optional.

        )")
        .def(py::init([](const std::string& filename,
                         const std::vector<sensor::sensor_info>& info,
                         const std::vector<std::string>& fields_to_write,
                         uint32_t chunk_size) {
                 return new osf::Writer(filename, info, fields_to_write,
                                        chunk_size);
             }),
             py::arg("filename"), py::arg("info"),
             py::arg("fields_to_write") = std::vector<std::string>{},
             py::arg("chunk_size") = 0,
             R"(
             Creates a `Writer` with specified ``chunk_size``.

             Default ``chunk_size`` is ``2MB``.

             Args:
                filename (str): The filename to output to.
                info (List[sensor_info]): The sensor info vector to use for a 
                    multi stream OSF file.
                fields_to_write (List[str]): The fields from scans to
                    actually save into the OSF. If not provided uses the fields from 
                    the first saved lidar scan for each stream. This parameter is optional.
                chunk_size (int): The chunksize to use for the OSF file, this arg
                    is optional.

        )")
        .def(
            "save",
            [](osf::Writer& writer, uint32_t stream_index,
               const LidarScan& scan) { writer.save(stream_index, scan); },
            py::arg("stream_index"), py::arg("scan"),
            R"(
               Save a lidar scan to the OSF file.

               Args:
                   stream_index (int): The index of the corrosponding 
                       sensor_info to use.
                   scan (LidarScan): The scan to save.

            )")
        .def(
            "save",
            [](osf::Writer& writer, uint32_t stream_index,
               const LidarScan& scan, uint64_t ts) {
                writer.save(stream_index, scan, ouster::osf::ts_t(ts));
            },
            py::arg("stream_index"), py::arg("scan"), py::arg("ts"),
            R"(
               Save a lidar scan to the OSF file.

               Args:
                   stream_index (int): The index of the corresponding 
                       sensor_info to use.
                   scan (LidarScan): The scan to save.
                   ts (int): The timestamp to index the scan with.
            )")
        .def(
            "save",
            [](osf::Writer& writer, const std::vector<LidarScan>& scans) {
                writer.save(scans);
            },
            py::arg("scan"),
            R"(
               Save a set of lidar scans to the OSF file.

               Args:
                   scans (List[LidarScan]): The scans to save. This will correspond
                       to the list of sensor_infos.

            )")
        .def(
            "set_metadata_id",
            [](osf::Writer& writer, const std::string& str) {
                return writer.set_metadata_id(str);
            },
            R"(
                 Set the metadata identifier string.
            )")
        .def(
            "metadata_id",
            [](osf::Writer& writer) { return writer.metadata_id(); },
            R"(
                 Return the metadata identifier string.

                 Returns (str):
                     The OSF metadata identifier string.
            )")
        .def(
            "filename", [](osf::Writer& writer) { return writer.filename(); },
            R"(
                 Return the osf file name.

                 Returns (str):
                     The OSF filename.
            )")
        .def(
            "add_metadata",
            [](osf::Writer& writer, py::object m) {
                uint32_t res = 0;
                if (py::hasattr(m, "type_id")) {
                    std::string type_id =
                        py::cast<std::string>(py::getattr(m, "type_id"));
                    osf::MetadataEntry* me = m.cast<osf::MetadataEntry*>();
                    res = writer.add_metadata(*me);
                }
                return res;
            },
            py::arg("m"), "Add `metadata entry` to a file")
        .def_property_readonly("meta_store", &osf::Writer::meta_store, R"(
                Returns the metadata store that gives an access to all
                *metadata entries* in the file.
        )")
        .def(
            "save_message",
            [](osf::Writer& writer, uint32_t stream_id, uint64_t receive_ts,
               uint64_t sensor_ts, py::array_t<uint8_t>& buf) {
                writer.save_message(stream_id, osf::ts_t{receive_ts},
                                    osf::ts_t{sensor_ts}, getvector(buf));
            },
            py::arg("stream_id"), py::arg("receive_ts"), py::arg("sensor_ts"),
            py::arg("buffer"), R"(
                 Low-level save message routine.

                 Directly saves the message `buffer` with `id` and `ts` (ns)
                 without any further checks.
            )")
        .def(
            "save_message",
            [](osf::Writer& writer, uint32_t stream_id, uint64_t receive_ts,
               uint64_t sensor_ts, py::buffer& buf) {
                writer.save_message(stream_id, osf::ts_t{receive_ts},
                                    osf::ts_t{sensor_ts}, getvector(buf));
            },
            py::arg("stream_id"), py::arg("receive_ts"), py::arg("sensor_ts"),
            py::arg("buffer"), R"(
                 Low-level save message routine.

                 Directly saves the message `buffer` with `id` and `ts` (ns)
                 without any further checks.
            )")
        .def(
            "add_sensor",
            [](osf::Writer& writer, const sensor::sensor_info& info,
               const std::vector<std::string>& fields_to_write) {
                return writer.add_sensor(info, fields_to_write);
            },
            py::arg("info"),
            py::arg("fields_to_write") = std::vector<std::string>{},
            R"(
               Add a sensor to the OSF file.

               Args:
                   info (sensor_info): Sensor to add.
                   fields_to_write (List[str]): The fields from scans to
                       actually save into the OSF. If not provided uses the fields from 
                       the first saved lidar scan for each stream. This parameter is optional.
               
               Returns (int):
                   The stream index to use to write scans to this sensor.

            )")
        .def("close", &osf::Writer::close,
             "Finish OSF file and flush everything on disk.")
        .def(
            "is_closed", [](osf::Writer& writer) { return writer.is_closed(); },
            R"(
                 Return the closed status of the writer.

                 Returns (bool):
                     The closed status of the writer.

            )")
        .def(
            "save",
            [](osf::Writer& writer, uint32_t stream_index,
               const LidarScan& scan) { writer.save(stream_index, scan); },
            py::arg("stream_index"), py::arg("scan"),
            R"(
               Save a lidar scan to the OSF file.

               Args:
                   stream_index (int): The index of the corrosponding 
                       sensor_info to use.
                   scan (LidarScan): The scan to save.

            )")
        .def(
            "save",
            [](osf::Writer& writer, const std::vector<LidarScan>& scans) {
                writer.save(scans);
            },
            py::arg("scan"),
            R"(
               Save a set of lidar scans to the OSF file.

               Args:
                   scans (List[LidarScan]): The scans to save. This will correspond
                       to the list of sensor_infos.

            )")
        .def(
            "sensor_info",
            [](osf::Writer& writer) { return writer.sensor_info(); },
            R"(
                 Return the sensor info list.

                 Returns (List[sensor_info]):
                     The sensor info list.

            )")
        .def(
            "sensor_info",
            [](osf::Writer& writer, uint32_t stream_index) {
                return writer.sensor_info(stream_index);
            },
            py::arg("stream_index"),
            R"(
                 Return the sensor info of the specifed stream_index.

                 Args:
                     stream_index (in): The index of the sensor to return
                                        info about.

                 Returns (sensor_info):
                     The correct sensor info

            )")
        .def(
            "sensor_info_count",
            [](osf::Writer& writer) { return writer.sensor_info_count(); },
            R"(
                 Return the number of sensor_info objects.

                 Returns (int):
                     The number of sensor_info objects.

            )")
        .def(
            "__enter__", [](osf::Writer* writer) { return writer; },
            R"(
                 Allow Writer to work within `with` blocks.
            )")
        .def(
            "__exit__",
            [](osf::Writer& writer, pybind11::object& /*exc_type*/,
               pybind11::object& /*exc_value*/,
               pybind11::object& /*traceback*/) {
                writer.close();

                return py::none();
            },
            R"(
                 Allow Writer to work within `with` blocks.
            )");

    m.def("slice_and_cast", &ouster::osf::slice_with_cast,
          py::arg("lidar_scan"), py::arg("field_types"),
          "Copies LidarScan with new field types");

    m.def(
        "slice_and_cast",
        [](const LidarScan& ls,
           const std::map<std::string, py::object>& field_types) {
            ouster::LidarScanFieldTypes ft{};
            for (const auto& f : field_types) {
                auto dtype = py::dtype::from_args(f.second);
                ft.push_back(FieldType(f.first, field_type_of_dtype(dtype)));
            }
            return ouster::osf::slice_with_cast(ls, ft);
        },
        py::arg("lidar_scan"), py::arg("field_types"),
        "Copies LidarScan with new field types");

    // TODO[pb]: (HACK) This is copied directly from _client.cpp to enable the
    // switch of the logger level in the compiled OSF SDK native module code.
    m.def(
        "init_logger",
        [](const std::string& log_level, const std::string& log_file_path,
           bool rotating, int max_size_in_bytes, int max_files) {
            return sensor::init_logger(log_level, log_file_path, rotating,
                                       max_size_in_bytes, max_files);
        },
        R"(
        Initializes and configures ouster_client logs. This method should be invoked
        only once before calling any other method from the library if the user wants
        to direct the library log statements to a different medium (other than
        console which is the default).

        Args:
            log_level Control the level of log messages outputed by the client.
                Valid options are (case-sensitive): "trace", "debug", "info", "warning",
                "error", "critical" and "off".
            log_file_path (str): Path to location where log files are stored. The
                path must be in a location that the process has write access to. If an empty
                string is provided then the logs will be directed to the console. When
                an empty string is passed then the rest of parameters are ignored.
            rotating (bool): Configure the log file with rotation, rotation rules are
                specified through the two following parameters max_size_in_bytes and
                max_files. If rotating is set to false the following parameters are ignored
            max_size_in_bytes (int): Maximum number of bytes to write to a rotating log
                file before starting a new file. ignored if rotating is set to False.
            max_files (int): Maximum number of rotating files to accumlate before
                re-using the first file. ignored if rotating is set to False.

        Returns:
            returns True on success, False otherwise.
        )",
        py::arg("log_level"), py::arg("log_file_path") = "",
        py::arg("rotating") = false, py::arg("max_size_in_bytes") = 0,
        py::arg("max_files") = 0);

    // TODO[pb]: (HACK) This is copied directly from _client.cpp to enable the
    // custom profile addition in the compiled OSF SDK native module code.
    m.def("add_custom_profile", &ouster::sensor::add_custom_profile);
}
